"""
PyTorch in action.
PyTorch 实现iris数据集的分类
"""

import torch
import matplotlib.pyplot as plt
import torch.nn.functional as F
from sklearn.datasets import load_iris
from torch.autograd import Variable
from torch.optim import SGD

# 动态地判断GPU是否可用，方便在不同类型的处理器上迁移
# GPU是否可用
use_cuda = torch.cuda.is_available()
print(f"use_cuda:{use_cuda}")

# 加载数据集
iris = load_iris()
# dict_keys('target_names','data','feature_names','DESCR','target')
# print(iris.keys())

# 数据预处理，包括从数据集里区分输入/输出，最后吧输入/输出数据封装成PyTorch期望的Variable格式
# 特征信息
x = iris['data']
# 目标分类
y = iris['target']
# print(x.shape)  # (150,4)
# print(y.shape)  # (150,)
# print(y)
x = torch.FloatTensor(x)
y = torch.LongTensor(y)
x, y = Variable(x), Variable(y)


# 神经网络模型定义，PyTorch中自定义的模型都需要继承Module，并重写forward方法完成前向计算过程
class Net(torch.nn.Module):
    # 初始化函数，接受自定义输入特征维数，隐含层特征维数，输出层特征维数
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__()
        # 一个线性隐含层
        self.hidden = torch.nn.Linear(n_feature, n_hidden)
        # 线性输出层
        self.predict = torch.nn.Linear(n_hidden, n_output)

    # 前向传播过程
    def forward(self, x):
        # x = F.sigmoid(self.hidden(x))
        x = torch.sigmoid(self.hidden(x))
        x = self.predict(x)
        out = F.log_softmax(x, dim=1)
        return out

# 网路实例化并打印查看网络结构
# iris中输入特征4维，隐含层和输出层可以自己选择
net=Net(n_feature=4,n_hidden=5,n_output=4)
# print(net)

# 判断GPU是否可用。如果GPU可用，就将训练数据和模型都放到GPU上，调用cuda()函数就可以把相应的模块
# 放到GPU上；相反，如果想放到GPU上，则调用.cpu()即可。注意：数据和网络到要在GPU上同步，否则程序
# 运行的时候会报错。

if use_cuda:
    x=x.cuda()
    y=y.cuda()
    net=net.cuda()

# 定义神经网络训练的优化器，并设置学习率为0.5
optimizer=SGD(net.parameters(),lr=0.5)
# 训练过程
px,py=[],[] #记录要绘制的数据
for i in range(1000):
    # 数据集传入网络前向计算
    prediction=net(x)
    # 计算loss
    loss=F.nll_loss(prediction,y)
    # 清除网络状态
    optimizer.zero_grad()
    # loss反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    # 在训练过程中打印每次迭代的损失情况
    # 打印并记录当前的index和loss
    print(i,"loss:",loss.data.item())
    px.append(i)
    py.append(loss.data.item())

    # 每10次迭代绘制训练动态
    # 动态画出loss走向结果在image/result_of_iris_multi-classfication.png
    plt.cla()
    plt.plot(px,py,'r-',lw=1)
    plt.text(0,0,'Loss=%.4f'%loss.data.item(),fontdict={'size':20,'color':'red'})
    plt.pause(0.1)

